package com.iscas.loadstorm.isloadstorm.util;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.iscas.loadstorm.isloadstorm.collector.CustomResultCollector;
import com.iscas.loadstorm.isloadstorm.domain.LoadApi;
import com.iscas.loadstorm.isloadstorm.domain.TestRun;
import org.apache.jmeter.config.Arguments;
import org.apache.jmeter.control.LoopController;
import org.apache.jmeter.engine.StandardJMeterEngine;
import org.apache.jmeter.protocol.http.control.Header;
import org.apache.jmeter.protocol.http.control.HeaderManager;
import org.apache.jmeter.protocol.http.sampler.HTTPSamplerProxy;
import org.apache.jmeter.reporters.Summariser;
import org.apache.jmeter.testelement.TestPlan;
import org.apache.jmeter.util.JMeterUtils;
import org.apache.jorphan.collections.HashTree;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;
import org.apache.jmeter.threads.ThreadGroup;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

@Component

public class ApiRunner {
    private static final Logger log = LogManager.getLogger(ApiRunner.class);

    @Autowired
    @Qualifier("customTaskExecutor")
    private ThreadPoolTaskExecutor taskExecutor;

    private static ConcurrentHashMap<String, TestRun> testRuns = new ConcurrentHashMap<>();
    private static final ObjectMapper objectMapper = new ObjectMapper();

    public String runApiTest(LoadApi loadApi, String jmeterHome) {
        String testId = UUID.randomUUID().toString();
        log.info("Starting API test with ID: {}, target service: {}", testId, loadApi.getTargetService());
        log.info("Load test configuration - Users: {}, Ramp-up time: {}, Load pattern: {}",
                loadApi.getUserCount(), loadApi.getRampUpTime(), loadApi.getLoadPattern());

        StandardJMeterEngine jmeterEngine = new StandardJMeterEngine();

        try {
            log.info("Initializing JMeter with home directory: {}", jmeterHome);
            initJMeter(jmeterHome);
            log.info("JMeter initialization completed successfully");

            log.info("Building test plan for test ID: {}", testId);
            HashTree testPlanTree = buildTestPlan(loadApi, testId, jmeterEngine);
            log.info("Test plan built successfully");

            log.info("Submitting test execution task to thread pool for test ID: {}", testId);
            taskExecutor.execute(() -> {
                try {
                    log.info("Configuring JMeter engine for test ID: {}", testId);
                    jmeterEngine.configure(testPlanTree);
                    log.info("Starting JMeter engine for test ID: {}", testId);
                    jmeterEngine.run();
                    log.info("JMeter engine completed execution for test ID: {}", testId);
                } catch (Exception e) {
                    log.error("Error executing JMeter test ID: {}", testId, e);
                }
            });

            testRuns.put(testId, new TestRun(testId, jmeterEngine, null));
            log.info("Test run record created and stored for test ID: {}", testId);
            return testId;
        } catch (Exception e) {
            log.error("Failed to start API test ID: {}. Error: {}", testId, e.getMessage(), e);
            throw new RuntimeException("Failed to start API test: " + e.getMessage());
        }
    }

    private static void initJMeter(String jmeterHome) {
        log.info("Setting JMeter home: {}", jmeterHome);
        JMeterUtils.setJMeterHome(jmeterHome);

        String propertiesPath = jmeterHome + "/bin/jmeter.properties";
        log.info("Loading JMeter properties from: {}", propertiesPath);
        JMeterUtils.loadJMeterProperties(propertiesPath);

        log.info("Initializing JMeter locale");
        JMeterUtils.initLocale();
    }

    private static HashTree buildTestPlan(LoadApi loadApi, String testId, StandardJMeterEngine jmeterEngine) {
        System.out.println("我进来啦");
        log.info("Building test plan components for test ID: {}", testId);

        TestPlan testPlan = new TestPlan("API Load Test Plan");
        log.info("Created test plan: {}", testPlan.getName());

        LoopController loopController = new LoopController();
        boolean isConstantLoad = loadApi.getLoadPattern().equals("0");
        loopController.setLoops(isConstantLoad ? -1 : 1);
        loopController.setContinueForever(isConstantLoad);
        loopController.initialize();
        log.info("Configured loop controller - Infinite loops: {}", isConstantLoad);

        ThreadGroup threadGroup = new ThreadGroup();
        threadGroup.setName("Load Test Thread Group");
        threadGroup.setNumThreads(loadApi.getUserCount());
        threadGroup.setRampUp(loadApi.getRampUpTime() != null ? loadApi.getRampUpTime() : 0);
        threadGroup.setSamplerController(loopController);
        log.info("Configured thread group - Threads: {}, Ramp-up: {}",
                threadGroup.getNumThreads(), threadGroup.getRampUp());

        HTTPSamplerProxy httpSampler = new HTTPSamplerProxy();
        httpSampler.setName("HTTP Request");
        httpSampler.setDomain(loadApi.getTargetService());
        httpSampler.setPath(loadApi.getRequestUrl());
        httpSampler.setMethod(loadApi.getHttpMethod());
        log.info("Configured HTTP sampler - Domain: {}, Path: {}, Method: {}",
                httpSampler.getDomain(), httpSampler.getPath(), httpSampler.getMethod());

        configureRequestParameters(httpSampler, loadApi);

        HeaderManager headerManager = new HeaderManager();
        if (loadApi.getCustomHeaders() != null) {
            log.info("Processing custom headers for test ID: {}", testId);
            String[] headers = loadApi.getCustomHeaders().split("\n");
            for (String header : headers) {
                String[] keyValue = header.split(":");
                if (keyValue.length == 2) {
                    headerManager.add(new Header(keyValue[0].trim(), keyValue[1].trim()));
                    log.trace("Added header - Key: {}, Value: {}", keyValue[0].trim(), keyValue[1].trim());
                }
            }
        }

        log.info("Building test plan tree hierarchy");
        HashTree testPlanTree = new HashTree();
        HashTree testPlanHashTree = testPlanTree.add(testPlan);
        HashTree threadGroupHashTree = testPlanHashTree.add(threadGroup);
        HashTree samplerHashTree = threadGroupHashTree.add(httpSampler);
        samplerHashTree.add(headerManager);

        TestRun testRun = new TestRun(testId, jmeterEngine, null);
        Summariser summariser = new Summariser("summary");
        String resultFile = "/Users/mj/Desktop/ISCS_doc/codedemo/jmeterdemo/test_results_java.csv";
        log.info("Configuring result collector with output file: {}", resultFile);

        CustomResultCollector resultCollector = new CustomResultCollector(testRun);
        resultCollector.setFilename(resultFile);
        testPlanHashTree.add(resultCollector);

        log.info("Test plan tree built successfully for test ID: {}", testId);
        return testPlanTree;
    }

    public static void configureRequestParameters(HTTPSamplerProxy httpSampler, LoadApi loadApi) {
        try {
            if (loadApi.getRequestBody() != null && !loadApi.getRequestBody().trim().isEmpty()) {
                log.info("Setting raw request body");
                httpSampler.addNonEncodedArgument("", loadApi.getRequestBody(), "");
                httpSampler.setPostBodyRaw(true);
                log.trace("Request body set: {}", loadApi.getRequestBody());
                return;
            }

            if (loadApi.getRequestParams() != null && !loadApi.getRequestParams().trim().isEmpty()) {
                log.info("Processing request parameters");
                String requestParams = loadApi.getRequestParams();
                log.trace("Raw request parameters: {}", requestParams);

                Map<String, Object> paramMap = objectMapper.readValue(requestParams, Map.class);
                log.info("Parsed {} parameters from request", paramMap.size());

                Arguments args = new Arguments();
                paramMap.forEach((key, value) -> {
                    args.addArgument(key, value.toString());
                    log.trace("Added parameter - Key: {}, Value: {}", key, value);
                });

                httpSampler.setArguments(args);
                log.info("Request parameters configured successfully");
            }
        } catch (Exception e) {
            log.error("Failed to configure request parameters", e);
            throw new RuntimeException("Failed to configure request parameters: " + e.getMessage());
        }
    }

    public static TestRun getTestRun(String testId) {
        log.info("Retrieving test run for ID: {}", testId);
        TestRun testRun = testRuns.get(testId);
        log.info("Test run {} found: {}", testId, testRun != null ? "yes" : "no");
        return testRun;
    }

    public static void removeTestRun(String testId) {
        log.info("Removing test run with ID: {}", testId);
        testRuns.remove(testId);
        log.info("Test run removed: {}", testId);
    }

    public static void stopTest(String testId) {
        log.info("Attempting to stop test with ID: {}", testId);
        TestRun testRun = testRuns.get(testId);
        if (testRun != null && testRun.getEngine() != null) {
            log.info("Stopping JMeter engine for test ID: {}", testId);
            testRun.getEngine().stopTest(true);
            removeTestRun(testId);
            log.info("Test stopped and removed successfully: {}", testId);
        } else {
            log.warn("Cannot stop test ID: {}. Test run or engine not found", testId);
        }
    }
}